import math

import numpy as np
import torch

torch.set_default_dtype(torch.float64)
torch.set_default_tensor_type(torch.DoubleTensor)

RESTART_NUMS = 1000
ROUND_NUMS = 1
EPOCHS_NUMS = 500

case_name = "eg5"
DIM = 3  # 系统输入维度
N_H = 1  # 隐藏层层数
D_H = 10  # 每个隐藏层包含的神经元个数
VERBOSE = True  # 是否打印输出
# VERBOSE = False  # 是否打印输出
FINE_TUNE = True  # false: 初始化参数, true: 微调参数，无需初始化参数
USE_CDINN = True  # false: nn, true: clamp(weight)
# USE_CDINN = False  # false: nn, true: clamp(weight)

alpha = 0.1  # B(f(x))-(1-alpha)B(x)
time_step = 0.05  # x_{t+1} = x_t + time_step * x'
piece_len = 0.3

BATCH_SIZE = 16  # mini_batch_train_size
# piece_num = [[16, 16, 16], [16, 16, 16], [32, 32, 32]]  # generate data, init/unsafe/domain
piece_num = [[8, 8, 8], [8, 8, 8], [16, 16, 16]]  # generate data, init/unsafe/domain
# piece_num = [[8, 8, 8], [8, 8, 8], [16, 16, 16]]  # generate data, init/unsafe/domain
cex_num = [[5, 5, 5], [5, 5, 5], [5, 5, 5]]  # add cex, init/unsafe/domain
cex_step = 0.001  # add_cex_step
cnt_num = 5
TOL_INIT = 0.0  # loss = relu(B(x) + tol)
TOL_UNSAFE = 0.0
TOL_DOMAIN = 0.0

DECAY_DOMAIN = 1  # p1 * init + p2 * unsafe + p3 * domain
DECAY_INIT = 1
DECAY_UNSAFE = 1

LOSS_OPT_FLAG = 1e-16  # epoch_loss < flag ==> train_success
TOL_DATA_GEN = 1e-16  # -1.8 + tol <= x0 <= 2 + tol; -0.5 + tol <= x0 <= 0.5 + tol

init_min = [0.5, 0.5, 0.5]
init_max = [0.9, 0.9, 0.9]
init_shape = 2

unsafe_min = [-0.9, -0.9, -0.9]
unsafe_max = [-0.5, -0.5, -0.5]
unsafe_shape = 2

domain_min = [-1, -1, -1]
domain_max = [1, 1, 1]


def is_unsafe(xt):
    return unsafe_min[0] <= xt[0] <= unsafe_max[0] and unsafe_min[1] <= xt[1] <= unsafe_max[1] and unsafe_min[2] <= xt[2] <= unsafe_max[2] \
        and (xt[0] + 0.7) * (xt[0] + 0.7) + (xt[1] + 0.7) * (xt[1] + 0.7) + (xt[2] + 0.7) * (xt[2] + 0.7) <= 0.2 * 0.2


def is_init(xt):
    return init_min[0] <= xt[0] <= init_max[0] and init_min[1] <= xt[1] <= init_max[1] and init_min[2] <= xt[2] <= init_max[2] \
        and (xt[0] - 0.7) * (xt[0] - 0.7) + (xt[1] - 0.7) * (xt[1] - 0.7) + (xt[2] - 0.7) * (xt[2] - 0.7) <= 0.2 * 0.2


def is_domain(xt):
    return domain_min[0] <= xt[0] <= domain_max[0] and domain_min[1] <= xt[1] <= domain_max[1] and domain_min[2] <= xt[2] <= domain_max[2]


def vector_field(x):
    def f(i, x):
        if i == 1:
            return time_step * (torch.cos(x[:, 2])) + x[:, 0]
        elif i == 2:
            return time_step * (torch.sin(x[:, 2])) + x[:, 1]
        elif i == 3:
            return time_step * (x[:, 2]) + x[:, 2]
        else:
            print("Vector function error!")
            exit()
    vf = torch.stack([f(i + 1, x) for i in range(DIM)], dim=1)
    return vf


# ############################################
# # for plotting
# ############################################
PLOT_EXP_B = np.array([7, 7, 7]) # sampling from domain for plotting the boundary of barrier using contour plot
PLOT_LEN_B = np.power(2, PLOT_EXP_B) # the number of samples for each dimension of domain, usually larger than superp.DATA_LEN_D
#
# PLOT_EXP_V = np.array([7, 7]) # sampling from domain for plotting the vector field
# PLOT_LEN_V = np.power(2, PLOT_EXP_V) # the number of samples for each dimension of domain, usually equal to superp.DATA_LEN_D
#
# PLOT_EXP_P = np.array([7, 7]) # sampling from domain for plotting the scattering sampling points, should be equal to superp.DATA_LEN_D
# PLOT_LEN_P = np.power(2, PLOT_EXP_P) # the number of samples for each dimension of domain
#
# PLOT_VEC_SCALE = None